#include "MSReader.h"
#include "Spectrum.h"
#include "MSObject.h"
#include <cmath>
#include <vector>

using namespace std;
using namespace MSToolkit;

#define GAUSSCONST 5.5451774444795623

typedef struct pep {
	char seq[64];
	double mass;
	int charge;
	int firstScan;
	int lastScan;
	float firstRT;
	float lastRT;
} pep;

typedef struct plotPoint {
	int charge;
	double mz;
	double ppm;
} plotPoint;

double calcFWHM(double mz,double res,int iType);
void centroid(Spectrum& s, Spectrum& out);
bool linearBestFit(vector<double>& p1, vector<double>& p2, double& slope, double& intercept);
void matchPeptides(Spectrum& s, vector<pep>& peps, vector<double>& vMZ, vector<double>& vPPM);
void readPeptides(char* fn, vector<pep>& vp);

int instrument;
double resolution;

int main(int argc, char* argv[]){

  cout << "RAWAccuracy, version 1.0" << endl;
  cout << "Copyright 2011-2013, Michael Hoopmann, Institute for Systems Biology" << endl;
	if(argc==1){
		cout << "USAGE: RAWAccuracy <Input file> <Peptide File> <Output File> [flags]" << endl;
		cout << "\tValid Input files: .RAW, .CMS1, .MS1, .mzXML, .mzML" << endl;
		cout << "\tValid Output files: .CMS1" << endl;
    cout << "\tPeptide File is tab delimited list containing the following columns:" << endl;
    cout << "\t\tPeptide Sequence" << endl;
    cout << "\t\tCharge" << endl;
    cout << "\t\tMass (zero charge)" << endl;
    cout << "\t\tFirst Retention Time" << endl;
    cout << "\t\tLast Retention Time" << endl;
		cout << "\tFLAGS:" << endl;
		cout << "\t\t-f\tData is from FT (assumed orbi data if not set)" << endl;
		cout << "\t\t-r <num>\tScan resolution (at m/z 400, default=60000)" << endl;
		exit(1);
	}

	MSReader r;
	MSReader w;
	MSObject o;

	Spectrum s;
	Spectrum c;
	Spectrum f;

	int i;

	bool bMatch;

	vector<double> vm;
	vector<double> vp;
	vector<pep> peps;

	double avg;
	double dif;
	double std;
	double intercept;
	double ppm;
	double slope;

	instrument=0;
	resolution=60000.0;

	if(argc>3){
		for(i=3;i<argc;i++){
			if(argv[i][0]!='-') continue;
			switch(argv[i][1]){
			case 'f':	instrument=3;	break;
			case 'r': resolution=atof(argv[i+1]);	break;
			default:
				cout << "Invalid flag!\n" << endl;
				exit(0);
				break;
			}
		}
	}

	FILE* fout=fopen("summary.txt","wt");

	//read in Peptides
	readPeptides(argv[2],peps);

	//read in first scan an set up output
	r.setFilter(MS1);
	r.readFile(argv[1],s);
	o.setHeader(r.getHeader());
	w.setCompression(true);
	w.writeFile(argv[3],cms1,o);

	//Iterate through all precursor scans
	while(s.getScanNumber()>0){
		centroid(s,c); //use centroids for speed and because it represents the true peak
		c.setRTime(s.getRTime());
		c.setScanNumber(s.getScanNumber());

		fprintf(fout,"%d\t",s.getScanNumber());

		//Find the peptides observed in the scan
		matchPeptides(c,peps,vm,vp);

		//If more than one peptide is found in the precursor scan, fit a function for mass accuracy
		if(vm.size()>1) {
			avg=0.0;
			dif=0.0;
			std=0.0;
			for(i=0;i<vp.size();i++) avg+=vp[i];
			avg/=vp.size();
			for(i=0;i<vp.size();i++) dif+=pow((vp[i]-avg),2);
			std=sqrt(dif/vp.size());
			//cout << "Mass accuracy of peptide precursors: " << avg << " +/- " << std << " Points: " << vp.size() << endl;
			fprintf(fout,"%d\t%.3lf\t%.3lf\t",vp.size(),avg,std);
			bMatch=linearBestFit(vm,vp,slope,intercept);
		} else {
			bMatch=false;
			fprintf(fout,"%d\tna\tna\tna\tna\tna\n",vp.size());
		}

		//Output the funtion to screen, then use it to apply a correction factor to the mass values
		//Then output the corrected spectra to file
		if(bMatch){			

			//cout << s.getScanNumber() << "\t";
			//cout << " y = " << slope << "x + " << intercept << endl;

			//apply correction to fixed spectrum object f
			f.clear();
			f.setScanNumber(s.getScanNumber());
			f.setScanNumber(s.getScanNumber(true),true);
			f.setRTime(s.getRTime());
			for(i=0;i<s.size();i++){
				f.add(s[i]);
				ppm = s[i].mz*slope+intercept;
				f[i].mz -= ppm/1000000.0*f[i].mz;
			}

			//check correction
			centroid(f,c); //use centroids for speed and because it represents the true peak
			c.setRTime(f.getRTime());
			c.setScanNumber(f.getScanNumber());
			matchPeptides(c,peps,vm,vp);

			avg=0.0;
			dif=0.0;
			std=0.0;
			for(i=0;i<vp.size();i++) avg+=vp[i];
			avg/=vp.size();
			for(i=0;i<vp.size();i++) dif+=pow((vp[i]-avg),2);
			std=sqrt(dif/vp.size());
			//cout << "Mass accuracy of peptide precursors after correction: " << avg << " +/- " << std << " Points: " << vp.size() << endl;
			fprintf(fout,"%d\t%.3lf\t%.3lf\n",vp.size(),avg,std);

			o.add(f);
		} else {
			//if we didn't have enough peptides to fit a function, keep the old scan
			//cout << "Not enough precursors or correction not possible." << endl;
			o.add(s);
		}
		if(o.size()>100) {
			w.appendFile(argv[3],o);
			o.clear();
		}
		r.readFile(NULL,s);
	}
	w.appendFile(argv[3],o);

	fclose(fout);
	return 0;
}

void matchPeptides(Spectrum& s, vector<pep>& peps, vector<double>& vMZ, vector<double>& vPPM){
	unsigned int i;
	
	int index;
	int j;
	
	double bestPPM;
	double mz;
	double ppm;

	bool bMatch;

	vMZ.clear();
	vPPM.clear();


	FILE* f=fopen("plot.txt","at");
	fprintf(f,"Scan:%d\n",s.getScanNumber());

	//Find the peptides observed in the scan
	for(i=0;i<peps.size();i++){
		if(s.getRTime()<peps[i].firstRT-0.25 || s.getRTime()>peps[i].lastRT+0.25) continue;
		mz = (peps[i].mass + peps[i].charge * 1.00727646677)/peps[i].charge;
		bMatch=false;
		bestPPM=100.0;
		for(j=0;j<s.size();j++){
			ppm=(s[j].mz-mz)/mz*1000000;
			if( fabs(ppm)<10.0 && fabs(ppm)<fabs(bestPPM) ){
				bMatch=true;
				bestPPM=ppm;
				index=j;
			}
		}
		if(bMatch){
			vMZ.push_back(s[index].mz);
			vPPM.push_back(bestPPM);
			fprintf(f,"%.4lf\t%.4lf\t%d\n",s[index].mz,bestPPM,peps[i].charge);
		}
	}
	fclose(f);
}

//Calculates the resolution (FWHM) of a peak
double calcFWHM(double mz,double res,int iType){
	double deltaM;
	switch(iType){
	case 0: //Orbitrap
		deltaM = mz * sqrt(mz) / (20*res);  //sqare root of 400
		break;
	case 1: //TOF
		deltaM = mz / res;
		break;
	case 2: //QIT
		deltaM = res;
		break;
	case 3: //FTICR
	default:
		deltaM = mz * mz / (400*res);
		break;
	}
	return deltaM;
}

//First derivative method, returns base peak intensity of the set
void centroid(Spectrum& s, Spectrum& out){
  int i,j;
  float maxIntensity;
  int bestPeak;
  bool bLastPos;

	int nextBest;
	double FWHM;
	Peak_T centroid;

	out.clear();

  bLastPos=false;
	for(i=0;i<s.size()-1;i++){

    if(s[i].intensity<s[i+1].intensity) {
      bLastPos=true;
      continue;
    } else {
      if(bLastPos){
				bLastPos=false;
	
        //Possible ways to improve this:
				//1. check FWHM - arg! what a computational disaster.
				//2. account for noise - another disaster.

				//find max and add peak
				maxIntensity=0;
				for(j=i;j<i+1;j++){
				  if (s[j].intensity>maxIntensity){
				    maxIntensity=s[j].intensity;
				    bestPeak = j;
				  }
				}

				//Best estimate of Gaussian centroid
				//Get 2nd highest point of peak
				if(bestPeak==s.size()){
					nextBest=bestPeak-1;
				} else if(s[bestPeak-1].intensity > s[bestPeak+1].intensity){
					nextBest=bestPeak-1;
				} else {
					nextBest=bestPeak+1;
				}

				//Get FWHM
				FWHM = calcFWHM(s[bestPeak].mz,resolution,instrument);

				//Calc centroid MZ (in three lines for easy reading)
				centroid.mz = pow(FWHM,2)*log(s[bestPeak].intensity/s[nextBest].intensity);
				centroid.mz /= GAUSSCONST*(s[bestPeak].mz-s[nextBest].mz);
				centroid.mz += (s[bestPeak].mz+s[nextBest].mz)/2;

				//Calc centroid intensity
				centroid.intensity=(float)(s[bestPeak].intensity/exp(-pow((s[bestPeak].mz-centroid.mz)/FWHM,2)*GAUSSCONST));

				//some peaks are funny shaped and have bad gaussian fit.
				//if error is more than 10%, keep existing intensity
				if( fabs((s[bestPeak].intensity - centroid.intensity) / centroid.intensity * 100) > 10 ||
            //not a good check for infinity
            centroid.intensity>999999999999.9 ||
            centroid.intensity < 0 ) {
					centroid.intensity=s[bestPeak].intensity;
				}

				//Hack until I put in mass ranges
				if(centroid.mz<0 || centroid.mz>2000) {
					//do nothing if invalid mz
				} else {
					out.add(centroid);
				}
			
      }

    }
  }

}

bool linearBestFit(vector<double>& p1, vector<double>& p2, double& slope, double& intercept){

	int i;
	double sumx=0.0;
	double sumy=0.0;
	double sumxy=0.0;
	double sumxx=0.0;
	double d;

	for(i=0;i<p1.size();i++){
    sumx+=p1[i];
    sumy+=p2[i];
    sumxy+=p1[i]*p2[i];
    sumxx+=p1[i]*p1[i];
  }

	d=(sumx*sumx - p1.size()*sumxx);
	if( d==0.0) return false;
	slope = (sumx*sumy - p1.size()*sumxy) / d;
  intercept = (sumx*sumxy - sumy*sumxx) / d;
	return true;
}

void readPeptides(char* fn,vector<pep>& vp){
	
	char str[1024];
	char* tok;
	pep p;

	vp.clear();
	
	FILE* f=fopen(fn,"rt");
	if(f==NULL) {
		cout << "Cannot open " << fn << endl;
		return;
	}

	while(!feof(f)){
		if(fgets(str,1024,f)==NULL) continue;
		tok=strtok(str,"\t\n");
		strcpy(p.seq,tok);
		tok=strtok(NULL,"\t\n");
		p.charge=atoi(tok);
		tok=strtok(NULL,"\t\n");
		p.mass=atof(tok);
		tok=strtok(NULL,"\t\n");
		p.firstRT=(float)atof(tok);
		tok=strtok(NULL,"\t\n");
		p.lastRT=(float)atof(tok);
		vp.push_back(p);
	}
	fclose(f);
}

